#ifdef WIN32

#include "ThreadAffinityWin.h"

bool ThreadAffinityWin::setAffinity(quint64 threadId, const QSet<int> &processorIndexs)
{
    // 检查处理器编号是否合法
    if (processorIndexs.isEmpty())
    {
        return false;
    }
    foreach (auto index, processorIndexs)
    {
        if (getActiveProcessorCount() - 1 < index ||
            0 > index)
        {
            return false;
        }
    }

    KAFFINITY mask = 0;
    foreach (auto index, processorIndexs)
    {
        int i = mapToGroupProcessorIndex(index);
        mask |= (0x01ULL << i);
    }

    int group = getGroupIndex(processorIndexs.values().first());
    return setAffinity(threadId, mask, group);
}

/**
 * @brief ThreadAffinityWin::setAffinity
 * 设置线程亲和性子函数
 * @param threadId 线程id
 * @param mask 掩码
 * @param group 处理器组序号
 * @return 设置结果
 */
bool ThreadAffinityWin::setAffinity(quint64 threadId, KAFFINITY mask, WORD group)
{
    // 线程id=>线程句柄
    HANDLE handle = OpenThread(THREAD_ALL_ACCESS, false, threadId);

#if _WIN32_WINNT >= 0x0601
    // 处理器组亲和性结构
    GROUP_AFFINITY group_affinity;
    group_affinity.Reserved[0] = 0;
    group_affinity.Reserved[1] = 0;
    group_affinity.Reserved[2] = 0;
    group_affinity.Mask = mask;
    group_affinity.Group = group;

    // 设置线程亲和性
    BOOL ret = SetThreadGroupAffinity(handle, &group_affinity, nullptr);
    return (ret != 0);
#else
    return false;
#endif
}

/**
 * @brief ThreadAffinityWin::getActiveProcessorCount
 * @return 返回系统上逻辑处理器总数量
 */
int ThreadAffinityWin::getActiveProcessorCount()
{
    int processorCount = 0;
    WORD groupCount = GetActiveProcessorGroupCount();
    for (WORD i = 0; i < groupCount; i++)
    {
        processorCount += GetActiveProcessorCount(i);
    }
    return processorCount;
}

/**
 * @brief ThreadAffinityWin::getGroupIndex
 * 获取处理器所属的处理器组序号
 * @param processorIndex 处理器序号
 * @return 处理器组序号
 */
int ThreadAffinityWin::getGroupIndex(int processorIndex)
{
    int count = processorIndex + 1;
    WORD groupCount = GetActiveProcessorGroupCount();
    for (WORD i = 0; i < groupCount; i++)
    {
        count = count - GetActiveProcessorCount(i);
        if (count <= 0)
        {
            return i;
        }
    }
    return -1; // processorIndex is invalid
}

/**
 * @brief ThreadAffinityWin::mapToGroupProcessorIndex
 * 处理器全局序号转换为处理器组范围内的序号。
 * @param processorIndex 处理器全局序号
 * @return 处理器组范围内的序号
 */
int ThreadAffinityWin::mapToGroupProcessorIndex(int processorIndex)
{
    int count = processorIndex + 1;
    WORD groupCount = GetActiveProcessorGroupCount();
    for (WORD i = 0; i < groupCount; i++)
    {
        count = count - GetActiveProcessorCount(i);
        if (count <= 0)
        {
            return (GetActiveProcessorCount(i) + count - 1);
        }
    }
    return -1; // processorIndex is invalid
}

#endif
